import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from k_means_双约束 import KMeans

data = pd.read_csv('../data/testData/double-constrain.csv')
# iris_types = ['SETOSA','VERSICOLOR','VIRGINICA']

x_axis = 'POINT_X'
y_axis = 'POINT_Y'

# plt.figure(figsize=(12,5))
# # '''figure(num=None, figsize=None, dpi=None, facecolor=None, edgecolor=None, frameon=True)'''
# # plt.subplot(1,2,1)
# # '''subplot(nrows,ncols,sharex,sharey,subplot_kw,**fig_kw)'''
#
# # for iris_type in iris_types:
# #     plt.scatter(data[x_axis][data['class']==iris_type],data[y_axis][data['class']==iris_type],label = iris_type)
# # plt.title('label known')
# # plt.legend()
# plt.scatter(data[x_axis][:],data[y_axis][:])
# plt.title('label unknown')
# plt.show()

num_examples = data.shape[0]
x_train = data.values.reshape(num_examples,13)

#指定好训练所需的参数
num_clusters = 5
max_iteritions =1000

k_means = KMeans(x_train,num_clusters)
centroids,closest_centroids_ids = k_means.train(max_iteritions)

# 可视化展示
# plt.subplot(1,2,2)
# plt.figure(figsize=(5,5))
# col = ['HotPink', 'Aqua', 'Chartreuse', 'yellow', 'LightSalmon']
# for i in range(num_clusters):
#     plt.scatter(data[x_axis][:],data[y_axis][:],  color=col[i])
#     plt.scatter([e[0] for e in closest_centroids_ids[i]], [e[1] for e in closest_centroids_ids[i]], color=col[i])
# plt.show()

#
plt.figure(figsize=(6,7))
# plt.subplot(1,2,1)
# for iris_type in iris_types:
#     plt.scatter(data[x_axis][data['class']==iris_type],data[y_axis][data['class']==iris_type],label = iris_type)
# plt.title('label known')
# plt.legend()
#
# plt.subplot(1,2,2)
for centroid_id, centroid in enumerate(centroids):
    current_examples_index = (closest_centroids_ids == centroid_id).flatten()
    plt.scatter(data[x_axis][current_examples_index],data[y_axis][current_examples_index],label = centroid_id)

for centroid_id, centroid in enumerate(centroids):
    plt.scatter(centroid[1],centroid[2],c='black',marker = 'x')
legend = plt.legend(["Partition 1", "Partition 2","Partition 3","Partition 4","Partition 5",])
plt.title('Double constrained spatial clustering')
plt.show()


